#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <sys/user.h>
#include <stdio.h>
#include <string.h>
#include <signal.h>
#include <sys/types.h>

#include "help_api.h"
#include "wrapper.h"

Wrapper::Wrapper(pid_t child)
{
    this->hd = shared_ptr<HardDiskController>(new HardDiskController(this));
    this->child = child;  
}

uint Wrapper::readWord(uint address)
{
    return ptrace(PTRACE_PEEKDATA, child, (uint*)address, 0);
}

void Wrapper::memRead(void* localDest, uint remoteSrc, uint size){
    uint wordSize = sizeof(uint);
    
    if(size % wordSize != 0)
    {
        error_print("wrapper.cpp: memRead was called with size size(%d) % wordSize(%d) != 0(%d)\n",
                    size, wordSize, size % wordSize);
        throw 0;
    }
    
    uint wordCount = size / wordSize;
    
    for(uint i=0;i<wordCount;i++){
        uint diff = (i*wordSize);
        uint data = readWord(remoteSrc + diff);
        *(uint*)((uint)localDest + diff) = data;
    }
}

void Wrapper::writeWord(uint address, uint data)
{
    ptrace(PTRACE_POKEDATA, child, address, data);
}

void Wrapper::memWrite(uint remoteDst, void* localSrc, uint size){
    uint wordSize = sizeof(uint);
    
    if(size % wordSize != 0)
    {
        error_print("wrapper.cpp: memWrite was called with size size(%d) % wordSize(%d) != 0(%d)\n",
                size, wordSize, size % wordSize);
        throw 0;
    }
    
    uint wordCount = size / wordSize;
    
    for(uint i=0;i<wordCount;i++){
        uint diff = (i*wordSize);
        writeWord(remoteDst + diff, *(uint*)((uint)localSrc + diff));
    }
    
}

Wrapper::~Wrapper()
{
    debug_print("~Wrapper()\n");
}

void Wrapper::handleSEGV(){
    unsigned char* opcode_bytes;
    unsigned char imm;
    unsigned int eax;
    unsigned char al;
    unsigned int edx;
    unsigned short dx;
    unsigned opcode;
    struct user_regs_struct regs;
    // read debuggee registers    
    ptrace(PTRACE_GETREGS, child, 0, &regs);
    // read the opcode that caused the fault (program counter) 
    opcode = ptrace(PTRACE_PEEKTEXT, child, regs.eip, 0);
    opcode_bytes = (unsigned char*)&opcode;

    // the following switch parsed all the possible variation of in/out opcodes on x86
    switch(opcode_bytes[0]){
        case 0xE7:
            // outl eax,imm8  
            imm = opcode_bytes[1];
            eax = regs.eax;
            //debug_print("wrapper.cpp: outl eax(0x%02X),imm8(0x%02X)\n", eax,imm);    
            hd->outp(imm, (uchar)eax);
            regs.eip += 2;
            break;
        case 0xE6:
            //outb al,imm8
            imm = opcode_bytes[1];
            eax = regs.eax;
            al = (unsigned char)eax;
            //debug_print("wrapper.cpp: outb al(0x%02X),imm8(0x%02X)\n", al,imm);
            hd->outp(imm, al);
            regs.eip += 2;
            break;
        case 0xEE:
            // outb al,dx
            eax = regs.eax;
            al = (unsigned char)eax;
            edx = regs.edx;
            dx = edx & 0xffff;
            //debug_print("wrapper.cpp: outb al(0x%02X),dx(0x%04X)\n", al, dx);          
            hd->outp(dx, al);
            regs.eip += 1;
            break;
        case 0xEF:
            // outl eax,dx
            eax = regs.eax;
            edx = regs.edx;
            dx = edx & 0xffff;
            //debug_print("wrapper.cpp: outl eax(0x%02X),dx(0x%04X)\n", eax, dx);          
            hd->outp(dx, (uchar)eax);
            regs.eip += 1;
            break;
        case 0xE5:
            // inl imm8,eax
            imm = opcode_bytes[1];
            eax = regs.eax;
            //debug_print("wrapper.cpp: inl imm(0x%02X),eax(0x%02X)\n", imm,eax);              
            regs.eax = hd->inp(imm);
            regs.eip += 2;
            break;
        case 0xE4:
            //inb imm8,al
            imm = opcode_bytes[1];
            eax = regs.eax;
            al = (unsigned char)eax;
            //debug_print("wrapper.cpp: inb imm(0x%02X),al(0x%02X)\n", imm,al);
            regs.eax = hd->inp(imm) & 0xff;
            regs.eip += 2;
            break;
        case 0xEC:
            // inb dx, al
            eax = regs.eax;
            al = (unsigned char)eax;
            edx = regs.edx;
            dx = edx & 0xffff;
            //debug_print("wrapper.cpp: inb dx(0x%02X),al(0x%04X)\n", dx, al);     
            regs.eax = hd->inp(dx) & 0xff;                    
            regs.eip += 1;
            break;
        case 0xED:
            // inl dx, eax
            eax = regs.eax;
            edx = regs.edx;
            dx = edx & 0xffff;
            //debug_print("wrapper.cpp: inl dx(0x%02X),eax(0x%04X)\n", dx, eax);
            regs.eax = hd->inp(dx);        
            regs.eip += 1;
            break;
        default:
            error_print("wrapper.cpp: unhandled opcode 0x%02X\n", opcode_bytes[0]);
            throw 0;
    };            
    // jump over the "faulty" in/out opcode and continue with execution
    ptrace(PTRACE_SETREGS, child, 0, &regs);
    ptrace(PTRACE_CONT, child, 0, 0);
}

void Wrapper::run()
{
    bool first_breakpoint = true;
    while(true){
        // wait for new signal from debuggee
        int wait_status;
        wait(&wait_status);
        
        if(WIFEXITED(wait_status)){
            debug_print("wrapper.cpp: child exited\n");
            return;
        }
        if(!WIFSTOPPED(wait_status))
            continue;
        switch(WSTOPSIG(wait_status))
        {
            case SIGSEGV:
                handleSEGV();
                break;
            case SIGTRAP:
                if(first_breakpoint)
                {
                    first_breakpoint = false;
                    if(ptrace(PTRACE_CONT, child, 0, 0) < 0) {
                        error_print("wrapper.cpp: ptrace failed\n");
                        throw 0;
                    }           
                    break;
                }else{
                    error_print("wrapper.cpp: more than one SIGTRAP, this is unexcpected\n");
                    throw 0;
                }
            default:
                error_print("wrapper.cpp: unhandled signal %s\n", strsignal(WSTOPSIG(wait_status)));
                throw 0;
        }
    }      
}